﻿/* FocusMeasure.cpp
 * 
 * Copyright (C) 2020, 2021 Yuan LI
 * email:18069211@qq.com
 * https://blog.csdn.net/liyuanbhu
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or (at
 * your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include "FocusMeasure.h"
#include "QImageUtils.h"
#include <QColor>

double roberts(const QImage &grayImage)
{
    if(grayImage.format() != QImage::Format_Grayscale8 || grayImage.format() != QImage::Format_Indexed8)
    {
        return -1.0;
    }
    double result = 0;
    int rows = grayImage.height();
    int cols = grayImage.width();
    for (int j = 0; j < rows - 1; j++)
    { // for all rows
        const uchar* current = grayImage.constScanLine(j); // current row
        const uchar* next = grayImage.constScanLine(j + 1); // next row
        for (int i = 0; i < cols - 1; i++)
        {
            // 1   0
            // 0  -1
            double Gx = current[i] - next[i + 1];
            // 0  -1
            // 1   0
            double Gy = next[i] - current[i + 1];
            result += Gx * Gx + Gy + Gy;
        }
    }
    return result/(rows * cols);
}

double sobel(const QImage &grayImage)
{
    if(grayImage.format() != QImage::Format_Grayscale8 || grayImage.format() != QImage::Format_Indexed8)
    {
        return -1.0;
    }
    double result = 0;
    int rows = grayImage.height();
    int cols = grayImage.width();
    for (int j = 1; j < rows - 1; j++)
    { // for all rows
        const uchar* previous = grayImage.constScanLine(j-1); // previous row
        const uchar* current = grayImage.constScanLine(j); // current row
        const uchar* next = grayImage.constScanLine(j+1); // next row

        for (int i = 1; i < cols - 1; i++)
        {
            // -1  0  1
            // -2  0  2
            // -1  0  1
            double Gx = previous[i + 1] + next[i + 1]
                    - previous[i - 1] - next[i - 1]
                    + 2 * current[i + 1] - 2 * current[i - 1];

            // -1 -2 -1
            //  0  0  0
            //  1  2  1
            double Gy = next[i - 1] + next[i + 1]
                    - previous[i - 1] - previous[i + 1]
                    + 2 * previous[i] - 2 * next[i];
            result += Gx * Gx + Gy + Gy;
        }
    }
    return result/(rows * cols);
}

double laplace(const QImage &grayImage)
{
    if(grayImage.format() != QImage::Format_Grayscale8 || grayImage.format() != QImage::Format_Indexed8)
    {
        return -1.0;
    }
    double result = 0;
    int rows = grayImage.height();
    int cols = grayImage.width();
    for (int j = 1; j < rows - 1; j++)
    { // for all rows
        const uchar* previous = grayImage.constScanLine(j-1); // previous row
        const uchar* current = grayImage.constScanLine(j); // current row
        const uchar* next = grayImage.constScanLine(j+1); // next row
        for (int i = 1; i < cols - 1; i++)
        {
            // 0  -1   0
            //-1   4  -1
            // 0  -1   0
            result +=  4 * current[i] - current[i - 1] - current[i + 1]
                    - previous[i] - next[i];
        }
    }
    return result/(rows * cols);
}

double fastRoberts(const QImage &grayImage)
{
    if(grayImage.format() != QImage::Format_Grayscale8 || grayImage.format() != QImage::Format_Indexed8)
    {
        return -1.0;
    }
    double result = 0;
    int rows = grayImage.height();
    int cols = grayImage.width();
    for (int j = 0; j < rows - 1; j += 2)
    { // for all rows
        const uchar* current = grayImage.constScanLine(j); // current row
        const uchar* next = grayImage.constScanLine(j + 1); // next row
        for (int i = 0; i < cols - 1; i += 2)
        {
            // 1   0
            // 0  -1
            double Gx = current[i] - next[i + 1];
            // 0  -1
            // 1   0
            double Gy = next[i] - current[i+1];
            result += Gx * Gx + Gy + Gy;
        }
    }
    return 4 * result/(rows * cols);
}

double fastSobel(const QImage &grayImage)
{
    if(grayImage.format() != QImage::Format_Grayscale8 || grayImage.format() != QImage::Format_Indexed8)
    {
        return -1.0;
    }
    double result = 0;
    int rows = grayImage.height();
    int cols = grayImage.width();
    for (int j = 1; j < rows - 1; j += 3)
    { // for all rows
        const uchar* previous = grayImage.constScanLine(j-1); // previous row
        const uchar* current = grayImage.constScanLine(j); // current row
        const uchar* next = grayImage.constScanLine(j+1); // next row

        for (int i = 1; i < cols - 1; i += 3)
        {
            // -1  0  1
            // -2  0  2
            // -1  0  1
            double Gx = previous[i + 1] + next[i + 1]
                    - previous[i - 1] - next[i - 1]
                    + 2 * current[i + 1] - 2 * current[i - 1];

            // -1 -2 -1
            //  0  0  0
            //  1  2  1
            double Gy = next[i - 1] + next[i + 1]
                    - previous[i - 1] - previous[i + 1]
                    + 2 * previous[i] - 2 * next[i];
            result += Gx * Gx + Gy + Gy;
        }
    }
    return 9*result/(rows * cols);

}

double fastLaplace(const QImage &grayImage)
{
    if(grayImage.format() != QImage::Format_Grayscale8 || grayImage.format() != QImage::Format_Indexed8)
    {
        return -1.0;
    }
    double result = 0;
    int rows = grayImage.height();
    int cols = grayImage.width();
    for (int j = 1; j < rows - 1; j += 3)
    { // for all rows
        const uchar* previous = grayImage.constScanLine(j-1); // previous row
        const uchar* current = grayImage.constScanLine(j); // current row
        const uchar* next = grayImage.constScanLine(j+1); // next row
        for (int i = 1; i < cols - 1; i += 3)
        {
            // 0  -1   0
            //-1   4  -1
            // 0  -1   0
            result +=  4 * current[i] - current[i - 1] - current[i + 1]
                    - previous[i] - next[i];
        }
    }
    return 9 * result/(rows * cols);
}

double discreteVariance(const QImage &grayImage)
{
    if(grayImage.format() != QImage::Format_Grayscale8 || grayImage.format() != QImage::Format_Indexed8)
    {
        return -1.0;
    }
    double mean = 0.0;
    double square = 0.0;
    double result = 0;
    int rows = grayImage.height();
    int cols = grayImage.width();
    for (int j = 0; j < rows - 1; j ++)
    { // for all rows
        const uchar* current = grayImage.constScanLine(j); // current row
        double mean_line = 0.0;
        double square_line = 0.0;
        for (int i = 0; i < cols - 1; i ++)
        {
            mean_line += current[j];
            square_line += current[j] * current[j];
        }
        mean_line /= cols;
        square_line /= cols;
        mean += mean_line;
        square += square_line;

    }
    mean /= rows;
    square /= rows;
    result = square - mean * mean;
    return result;
}

FocusMeasure::FocusMeasure()
{
    setMeasureMethod(ROBERTS);
    setZoomFactor(0);
}

double FocusMeasure::score(const QImage &image) const
{
    QImage grayImage = preprocess(image);
    return measure(grayImage);
}

double FocusMeasure::scoreAndRecord(const QImage &image, bool *better)
{
    double s = score(image);
    if(s > m_bestScore)
    {
        m_bestScore = s;
        m_bestImage = image;
        if(better != nullptr)
        {
            *better = true;
        }
    }
    else
    {
        if(better != nullptr)
        {
            *better = false;
        }
    }
    return s;
}

void FocusMeasure::reset()
{
    m_bestImage = QImage();
    m_bestScore = 0.0;
}

void FocusMeasure::setMeasureMethod(MeasureMethod method)
{
    switch(method)
    {
    case ROBERTS:
        m_measureFunc = roberts;
        break;
    case FAST_ROBERTS:
        m_measureFunc = fastRoberts;
        break;
    case SOBEL:
        m_measureFunc = sobel;
        break;
    case FAST_SOBEL:
        m_measureFunc = fastSobel;
        break;
    case LAPLACE:
        m_measureFunc = laplace;
        break;
    case FAST_LAPLACE:
        m_measureFunc = fastLaplace;
        break;
    case DISCRETE_VARIANCE:
        m_measureFunc = discreteVariance;
        break;
    default:
        m_measureFunc = roberts;
    }
}

QImage FocusMeasure::preprocess(const QImage &image) const
{
    QImage grayImage;
    if(m_roi.isValid())
    {
        QImage roiImage = roi(image, m_roi);
        grayImage = convertToGray(roiImage);
    }
    else
    {
        grayImage = convertToGray(image);
    }
    if(m_zoomFactor > 1)
    {
        int w = grayImage.width() / m_zoomFactor;
        grayImage = grayImage.scaledToWidth(w);
    }
    return grayImage;
}
